"""
训练最优的 K 值并保存 KNN 模型。

流程：
1) 加载 sklearn digits(8x8) 数据集
2) 训练/验证切分
3) 遍历 k=1..40 训练 KNN，记录测试集准确率
4) 保存最优模型到当前目录 best_knn.pkl
5) 打印最佳准确率和 k 值
"""

# TODO: 导入必要的库和模块
from __future__ import annotations

import pickle
from pathlib import Path
from typing import List, Tuple

import numpy as np
from sklearn import datasets
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import matplotlib
matplotlib.use("Agg")  # 后端设为非交互，避免弹窗
import matplotlib.pyplot as plt

# TODO: 加载数字数据集
def load_data() -> Tuple[np.ndarray, np.ndarray]:
	digits = datasets.load_digits()
	X, y = digits.data, digits.target  # X shape: (n_samples, 64), values in [0..16]
	return X, y

# TODO: 将数据集划分为训练集和测试集
def split_data(X: np.ndarray, y: np.ndarray, test_size: float = 0.2, random_state: int = 42):
	return train_test_split(
		X,
		y,
		test_size=test_size,
		random_state=random_state,
		stratify=y,
	)

# TODO: 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
def train_and_select_best(
	X_train: np.ndarray,
	y_train: np.ndarray,
	X_test: np.ndarray,
	y_test: np.ndarray,
	k_min: int = 1,
	k_max: int = 40,
):
	best_acc: float = -1.0
	best_k: int | None = None
	best_model: KNeighborsClassifier | None = None

	# TODO: 初始化一个列表以存储每个k值的准确率
	k_accuracies: List[float] = []

	# TODO: 尝试从1到40的k值，对于每个k值，训练knn模型，保存最佳准确率，k值和knn模型
	for k in range(k_min, k_max + 1):
		model = KNeighborsClassifier(n_neighbors=k)
		model.fit(X_train, y_train)
		y_pred = model.predict(X_test)
		acc = accuracy_score(y_test, y_pred)
		k_accuracies.append(acc)

		if acc > best_acc:
			best_acc = acc
			best_k = k
			best_model = model

	assert best_model is not None and best_k is not None
	return best_model, best_k, best_acc, k_accuracies

def main():
	X, y = load_data()
	X_train, X_test, y_train, y_test = split_data(X, y)

	best_model, best_k, best_acc, k_accuracies = train_and_select_best(
		X_train, y_train, X_test, y_test, k_min=1, k_max=40
	)

	# TODO: 将最佳KNN模型保存到二进制文件
	out_path = Path(__file__).resolve().parent / "best_knn.pkl"
	with out_path.open("wb") as f:
		pickle.dump({
			"model": best_model,
			"best_k": best_k,
			"best_acc": best_acc,
		}, f)

	# TODO: 打印最佳准确率和相应的k值
	print(f"Best accuracy: {best_acc:.4f} with k={best_k}")
	# 可选：打印部分 k 的准确率概览
	print("Top-5 k by accuracy:")
	top5 = sorted([(i + 1, acc) for i, acc in enumerate(k_accuracies)], key=lambda x: x[1], reverse=True)[:5]
	for k, acc in top5:
		print(f"  k={k:2d} -> acc={acc:.4f}")

	# 生成准确率曲线 PDF
	ks = list(range(1, len(k_accuracies) + 1))
	plt.figure(figsize=(8, 5))
	plt.plot(ks, k_accuracies, marker='o', label='Accuracy')
	plt.xlabel('k (n_neighbors)')
	plt.ylabel('Accuracy')
	plt.title('KNN Accuracy vs k on sklearn digits')
	# 垂直红线标注最佳 k
	plt.axvline(x=best_k, color='red', linestyle='--', linewidth=1.5, label=f'Best k={best_k}')
	# 标注交点文字
	best_y = k_accuracies[best_k - 1]
	plt.scatter([best_k], [best_y], color='red')
	plt.text(best_k + 0.5, best_y, f'k={best_k}, acc={best_y:.4f}', color='red')
	plt.legend()
	plt.grid(True, linestyle='--', alpha=0.3)

	pdf_path = Path(__file__).resolve().parent / 'accuracy_plot.pdf'
	plt.tight_layout()
	plt.savefig(pdf_path, format='pdf')
	plt.close()


if __name__ == "__main__":
	main()